from rich.console import Console
from rich.progress import Progress
import time, os, shutil
import struct, base64, ast
import zipfile

__version__ = '0.0.1'
console = Console()

'''
exam = {
    'f':{'bname':'size', 'bname2':'size2'},
    'd':{'bname':exam}
}
'''

def Hsize(size):
    def nchan(integer, remainder, level):
        if integer >= 1024:
            remainder = integer % 1024
            integer //= 1024
            level += 1
            return nchan(integer, remainder, level)
        else:
            return integer, remainder, level

    units = ['B', 'KiB', 'MiB', 'GiB', 'TiB', 'PiB']
    integer, remainder, level = nchan(size, 0, 0)
    if level+1 > len(units):
        level = -1
    return ( '{}.{:>03d} {}'.format(integer, remainder, units[level]) )

def B64en(data):
    data = data.encode('utf-8')
    data = base64.b64encode(data).decode('utf-8')
    return data

def B64de(data):
    data = data.encode('utf-8')
    data = base64.b64decode(data).decode('utf-8')
    return data

def CreatTree(path, con=console):
    if not os.path.isdir(path):
        con.print('[bold red]Target Not Exist.')
        return
    td = {'f':{},'d':{}}
    size = 0
    try:
        for i in os.listdir(path):
            p = os.path.join(path, i)
            i = B64en(i)
            if os.path.isfile(p):
                td['f'][i] = os.path.getsize(p)
                size += os.path.getsize(p)
                continue
            if os.path.isdir(p):
                td['d'][i] = CreatTree(p)
                continue
    except:return {}
    if td['f'] == {}:del td['f']
    if td['d'] == {}:del td['d']
    return td

def IndexTree(path, data, con=console):
    if not os.path.isdir(path):
        con.print('[bold red]Target Not Exist.')
        return
    if not data.get('d'):return
    for d in data['d']:
        dn = B64de(d)
        tp = os.path.join(path, dn)
        try:os.mkdir(tp)
        except:pass
        IndexTree(tp, data['d'][d])

def ReadFileTree(path, data, file, bar, con=console):
    progress, task = bar
    if not os.path.isdir(path):
        con.print('[bold red]Target Not Exist.')
        return
    if data.get('f'):
        for d in data['f']:
            size = data['f'][d]
            dn = B64de(d)
            tp = os.path.join(path, dn)
            with open(tp, 'rb') as f:
                while size > 0:
                    if size - 1024 > 0:
                        file.write(f.read(1024))
                        size -= 1024
                        progress.update(task, advance=1024)
                    else:
                        file.write(f.read(size))
                        progress.update(task, advance=size)
                        size = 0
    if data.get('d'):
        for d in data['d']:
            dn = B64de(d)
            tp = os.path.join(path, dn)
            ReadFileTree(tp, data['d'][d], file, bar)

def WriteFileTree(path, data, file, bar, con=console):
    progress, task = bar
    if not os.path.isdir(path):
        con.print('[bold red]Target Not Exist.')
        return
    if data.get('f'):
        for d in data['f']:
            size = data['f'][d]
            dn = B64de(d)
            tp = os.path.join(path, dn)
            with open(tp, 'wb+') as f:
                while size > 0:
                    if size - 1024 > 0:
                        f.write(file.read(1024))
                        size -= 1024
                        progress.update(task, advance=1024)
                    else:
                        f.write(file.read(size))
                        progress.update(task, advance=size)
                        size = 0
    if data.get('d'):
        for d in data['d']:
            dn = B64de(d)
            tp = os.path.join(path, dn)
            WriteFileTree(tp, data['d'][d], file, bar)

def TotalSize(data):
    size = 0
    if data.get('f'):
        for d in data['f']:
            size += data['f'][d]
    if data.get('d'):
        for d in data['d']:
            size += TotalSize(data['d'][d])
    return size

def Utif(path, fd, con=console):
    if not os.path.isdir(path):
        con.print('[bold red]Target Not Exist.')
        return
    with console.status('[bold blue]Indexing Tree...[/]'):
        data = CreatTree(path)
        size = TotalSize(data)
        lis = str(data).encode()
        length = len(lis)
        header = struct.pack('i', length)
        alls = bytearray(header+lis)
    con.print('[bold green]Tree Has Been Created.')
    con.print('[bold]Header Total : [bold]'+str(length))
    con.print('[bold]Files Total : [bold blue]'+Hsize(size))
    with Progress() as progress, open(fd+'p', 'wb+') as f:
        task = progress.add_task('[bold blue]Copy into.', total=size+1)
        f.write(alls)
        ReadFileTree(path, data, f, (progress, task))
        progress.update(task, advance=1)
    with console.status('[bold blue]Compressing File...[/]'):
        zip_file = zipfile.ZipFile(fd+'.zp', 'w')
        zip_file.write(fd+'p', compress_type=zipfile.ZIP_DEFLATED)
        zip_file.close()
        try: os.remove(fd)
        except:pass
        os.remove(fd+'p')
        os.rename(fd+'.zp', fd)
    con.print('[bold]Compressed into [bold blue]{}[bold].'.format(Hsize(os.path.getsize(fd))))
    con.print('[bold green]Successfully.')

def Utid(path, fd, con=console):
    if not os.path.isdir(path):
        os.mkdir(path)
        Utid(path, fd, con=con)
        return
    if not os.path.isfile(fd):
        con.print('[bold red]File Not Exist.')
        return
    tmpath, tmpname = os.path.split(fd)
    with console.status('[bold blue]Uncompressing File...[/]'):
        zip_file = zipfile.ZipFile(fd)
        zip_extract = zip_file.extractall(tmpath)
    con.print('[bold green]File Has Been Uncompressed.')
    with console.status('[bold blue]Reading File...[/]'), open(fd+'p', 'rb') as f:
        length = struct.unpack('i', f.read(4))[0]
        data = f.read(length).decode()
        data = ast.literal_eval(data)
        size = TotalSize(data)
        IndexTree(path, data)
    con.print('[bold]Files Total : [bold blue]'+Hsize(size))
    with Progress() as progress, open(fd+'p', 'rb') as f:
        f.seek(length+4)
        task = progress.add_task('[bold blue]Copy out.', total=size+1)
        WriteFileTree(path, data, f, (progress, task))
        progress.update(task, advance=1)
    with console.status('[bold blue]Deleting Temp File...[/]'):
        os.remove(fd+'p')
    con.print('[bold green]Successfully.')

def Comp7Z():
    pass